#! /bin/python3
import rospy
import numpy as np
import cv2
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image
from fyt_msg.msg import CvRect,AllTargets,DetectorParam

class Config_param:

    input_size = []
    compute_device = []
    path_model = []
    class_names_path = []
    class_names = []
    get_param = False #标志位，订阅并处理图像要在获取到参数之后

    def __init__(self):

        self.detect_param_sub = rospy.Subscriber("config_detector", DetectorParam, self.detectParamCallback)
        
    def detectParamCallback(self,data):
    
        self.get_param = False # 参数动态变化
        self.DEBUG = data.debug
        
        if self.DEBUG:
            print('<MSG>getting detect param')

        self.input_size = int(data.input_size)
        if data.device not in ['nx', 'manifold_g', 'manifold_c', 'nuc'] and self.DEBUG:
            print('<error> compute_device set error, set it as NX defaultly')
        else:
            self.compute_device = data.device
        self.path_model = data.model_path
        self.class_names_path = data.class_name_path

        file = open(self.class_names_path)
        self.class_names = [] #注意有self，否则不会清空
        for line in file.readlines():  
            curLine = line.strip().split(" ")      
            self.class_names.append(curLine[0]) #取每行第一个单词
        if self.DEBUG:
            print('<MSG>class_names: {}\n'.format(self.class_names))

        self.confThr = float(data.confThr)
        self.nmsThr = float(data.nmsThr)
        self.cal_time = data.calc_time

        self.get_param = True

class Detector:

    Detect_param = Config_param()
    #while not Detect_param.get_param: 会一直陷在死循环
    #    print(Detect_param.get_param)
    images = []
    #net = cv2.dnn.readNetFromONNX(Detect_param.path_model)
    net = None

    def __init__(self):
        self.fps = 0
        self.bridge = CvBridge()
        self.image_sub = rospy.Subscriber("raw_img", Image, self.imageColorCallback)
        self.AllTargets_pub = rospy.Publisher('/all_targets',AllTargets,queue_size=20)

    def imageColorCallback(self,data):
    
        if not self.Detect_param.get_param: # 判断Detect_param类参数是否未被初始化
            #self.net = None动态加载改变模型太影响性能
            return
        elif self.net == None: # 再判断Detector参数是否被初始化，否就初始化本类参数
            self.net = cv2.dnn.readNetFromONNX(self.Detect_param.path_model)
            if self.Detect_param.compute_device in ['nx', 'manifold_g']:
                #self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
                #self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA_FP16)
                pass
            elif self.Detect_param.compute_device == 'manifold_c':
                pass
            else:
                pass
        else: # 两个类的参数都被初始化完毕
            pass
    
        try:
            cv_images = self.bridge.imgmsg_to_cv2(data, "bgr8")
        except CvBridgeError as err:
            if self.Detect_param.DEBUG:
                print(err)
        
        if self.Detect_param.cal_time:
            start = cv2.getTickCount()
        image = Armor_Detector.pretreat(cv_images)
        output = Armor_Detector.postprocess(image)
        if self.Detect_param.cal_time:
            end = cv2.getTickCount()
            self.fps = 1000 * (end - start) / cv2.getTickFrequency()
        
        if self.Detect_param.DEBUG:
            Armor_Detector.draw(output)
        else:
            cv2.destroyAllWindows()
        
        CvRect_msg = CvRect()
        AllTargets_msg = AllTargets()
        
        '''
        if len(output) == 0:
            
            CvRect_msg.x = 0
            CvRect_msg.y = 0
            CvRect_msg.height = 0
            CvRect_msg.width = 0
            AllTargets_msg.rect.append(CvRect_msg) # 使用append的方法会使得其长度加一
            print(len(AllTargets_msg.rect))
            AllTargets_msg.id.append(0)
            
            AllTargets_msg.len = 0
            self.AllTargets_pub.publish(AllTargets_msg)
        '''
        
        AllTargets_msg.len = len(output)
        for (index, i) in enumerate(output):
            CvRect_msg.x = int(i[0])
            CvRect_msg.y = int(i[1])
            CvRect_msg.height = int(i[3] - i[1])
            CvRect_msg.width = int(i[2] - i[0])
            if index > 127:
                break
            AllTargets_msg.rect[index] = CvRect_msg
            AllTargets_msg.id[index] = int(i[-1])

        self.AllTargets_pub.publish(AllTargets_msg)

    def pretreat(self, images):
        #image_path="./test.jpg" 直接读取本功能包中imgs文件夹下图片测试接口留存
        #images=cv2.imread(image_path)
        self.images = images
        image_shape=images.shape[:2]
        image=images.copy()
        new_unpad = (self.Detect_param.input_size, self.Detect_param.input_size)
        image = cv2.resize(image,new_unpad, interpolation=cv2.INTER_LINEAR)
        image = image.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB  # should be setted
        image = np.ascontiguousarray(image)
        image=image[None]
        image=image/255.0

        return image

    def postprocess(self, image):
        image_shape=self.images.shape[:2]
        new_unpad = (self.Detect_param.input_size, self.Detect_param.input_size)
        self.net.setInput(image)
        pred = self.net.forward()
        
        #解码
        #[x, y, w, h] to [x1, y1, x2, y2]
        shape_boxes=np.zeros(pred[:,:,:4].shape)
        shape_boxes[:,:,0]=pred[:,:,0]-pred[:,:,2]/2
        shape_boxes[:,:,1]=pred[:,:,1]-pred[:,:,3]/2
        shape_boxes[:,:,2]=pred[:,:,0]+pred[:,:,2]/2
        shape_boxes[:,:,3]=pred[:,:,1]+pred[:,:,3]/2
        pred[:,:,:4]=shape_boxes

        output=[]
        for i in range(pred.shape[0]):
            prediction=pred[i]
            mask=(prediction[:,4] > self.Detect_param.confThr)
            #初步筛选置信度大于0.5
            prediction=prediction[mask]
            class_conf,class_pred=np.expand_dims(np.max(prediction[:,5:],1),-1),np.expand_dims(np.argmax(prediction[:,5:],1),-1)
            #种类得分
            detections=np.concatenate((prediction[:,:5],class_conf,class_pred),1)

            unique_class=np.unique(detections[:,-1])
            if self.Detect_param.DEBUG:
                print("<MSG>object counts: {}".format(len(unique_class)))
            best_box=[]
            #非极大抑制
            for c in unique_class:
                cls_mask=detections[:,-1]==c
                detection=detections[cls_mask]
                #同一种类
                arg_sort=np.argsort(detection[:,4])[::-1]
                detection=detection[arg_sort]
                #根据得分从大到小进行排序
                while np.shape(detection)[0] >0:
                    best_box.append(detection[0])
                    if len(detection) ==1:
                        break
                    #计算一个框和同种类后面所有的iou
                    ious = self.iou(best_box[-1],detection[1:])
                    detection=detection[1:][ious < self.Detect_param.nmsThr]
                    #   iou>nms阈值的被剔除 
            output.append(best_box)
        output=np.array(output)
        output=np.squeeze(output,axis=0)

        #还原原图比例
        #print(output)
        #exit(0)
        if len(output) != 0:
            output[:,0]*=image_shape[1]/new_unpad[0]
            output[:,1]*=image_shape[0]/new_unpad[1]
            output[:,2]*=image_shape[1]/new_unpad[0]
            output[:,3]*=image_shape[0]/new_unpad[1]

        return output

    def iou(self, b1, B2):
        inter_rect_x1=np.maximum(b1[0],B2[:,0])
        inter_rect_y1=np.maximum(b1[1],B2[:,1])
        inter_rect_x2=np.minimum(b1[2],B2[:,2])
        inter_rect_y2=np.minimum(b1[3],B2[:,3])

        inter_area=np.maximum(inter_rect_x2-inter_rect_x1,0)*np.maximum(inter_rect_y2-inter_rect_y1,0)
        area_b1=(b1[2]-b1[0])*(b1[3]-b1[1])
        area_b2=(B2[:,2]-B2[:,0])*(B2[:,3]-B2[:,1])
        iou=inter_area/np.maximum((area_b1+area_b2-inter_area),1e-6)

        return iou

    def draw(self, output): #画图
        for i in output:
            self.images=cv2.rectangle(self.images,(int(i[0]),int(i[1])),(int(i[2]),int(i[3])),(255,0,0),2)
            self.images=cv2.putText(self.images,"%s:%.2f"%(self.Detect_param.class_names[int(i[-1])],i[-2]),(int(i[0]),int(i[1])),cv2.FONT_HERSHEY_COMPLEX,1,(255,0,0),1)
        if self.Detect_param.cal_time:
            cv2.putText(self.images, "FPS: " + str(int(self.fps)), (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 3)
        cv2.imshow("CNN",self.images)
        cv2.waitKey(1)

if __name__ == '__main__':
    Armor_Detector = Detector()
    rospy.init_node('armor_detect_node')
    try:
        rospy.spin()
    except KeyboardInterrupt:
        print("Shutting down")
    cv2.destoryAllWindows()
